package cn.slipi.common.context;

import cn.slipi.common.constants.TraceConstants;

import java.util.HashMap;
import java.util.Map;

/**
 * @author lee
 */
public class ThreadContext {
    private static final ThreadLocal<Map<String, Object>> CTX_HOLDER = new ThreadLocal<>();

    public ThreadContext() {
    }

    public static Map<String, Object> get() {
        return CTX_HOLDER.get();
    }

    protected static Object getAttribute(String attributeName) {
        Map<String, Object> ctx = CTX_HOLDER.get();
        return ctx == null ? null : ctx.get(attributeName);
    }

    protected static void setAttribute(String attributeName, Object value) {
        Map<String, Object> ctx = CTX_HOLDER.get();
        if (ctx == null) {
            ctx = new HashMap<>();
            CTX_HOLDER.set(ctx);
        }

        ctx.put(attributeName, value);
    }

    public static String getTraceId() {
        return (String) getAttribute(TraceConstants.TRACE_ID);
    }

    public static void putTraceId(String traceId) {
        setAttribute(TraceConstants.TRACE_ID, traceId);
    }

    public static void removeAll() {
        CTX_HOLDER.remove();
    }
}